"""
Gemma 3 270M 模型加载器 - 使用 PyTorch + 编译优化（类似 LiteRT）
专为小尺寸模型优化，提供最佳推理性能
注意：Gemma-3 目前没有 TensorFlow 实现，使用 PyTorch + torch.compile 优化
"""
import os
import logging
from typing import Optional
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

logger = logging.getLogger(__name__)


class Gemma3ModelLoader:
    """Gemma 3 270M模型加载器 - PyTorch + 编译优化版"""

    def __init__(self, model_id: str = "google/gemma-3-270m", cache_dir: Optional[str] = None, local_model_path: Optional[str] = None):
        self.model_id = model_id
        self.cache_dir = cache_dir or os.getenv("MODEL_CACHE_DIR", "/app/models")
        self.local_model_path = local_model_path or os.getenv("LOCAL_MODEL_PATH", None)
        self.device = "cpu"  # CPU 推理
        self.model = None
        self.tokenizer = None
        self._compiled_model = None  # 编译后的模型

    def _ensure_chat_template(self):
        """
        确保 tokenizer.chat_template 存在
        如果模型是 base 模型，则自动注入 Gemma 通用 chat 模板
        """
        if getattr(self.tokenizer, "chat_template", None):
            logger.info("chat_template 已存在，无需注入。")
            return

        logger.warning("chat_template 缺失！自动注入 Gemma 通用对话模板。")

        # 通用 Gemma/Llama/Mistral 风格模板（兼容大多数 instruct 模型）
        default_template = """{% for message in messages -%}
<|{{ message['role'] }}|>
{{ message['content'] }}
{% endfor %}<|assistant|>"""

        self.tokenizer.chat_template = default_template
        logger.info("chat_template 注入成功。")

    def load_model(self):
        """加载模型 - 使用 PyTorch + torch.compile 优化（类似 LiteRT）"""
        try:
            logger.info("使用 PyTorch + 编译优化加载 Gemma-3-270M...")
            logger.info("注意：Gemma-3 目前没有 TensorFlow 实现，使用 PyTorch + torch.compile 优化")

            # 检查是否有本地模型路径
            if self.local_model_path and os.path.exists(self.local_model_path):
                logger.info(f"从本地路径加载模型: {self.local_model_path}")
                model_path = self.local_model_path
            else:
                logger.info(f"从HuggingFace加载模型: {self.model_id}")
                model_path = self.model_id

            # 加载tokenizer
            logger.info("加载 tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                model_path,
                cache_dir=self.cache_dir,
                trust_remote_code=True,
                local_files_only=self.local_model_path is not None and os.path.exists(self.local_model_path)
            )

            # 修复 chat_template 缺失问题
            self._ensure_chat_template()

            # 设置 pad_token
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            # 加载 PyTorch 模型
            logger.info("加载 PyTorch 模型...")
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path,
                cache_dir=self.cache_dir,
                trust_remote_code=True,
                local_files_only=self.local_model_path is not None and os.path.exists(self.local_model_path),
                torch_dtype=torch.float32,  # CPU 使用 float32
                low_cpu_mem_usage=True
            )

            # 移动到 CPU
            self.model = self.model.to(self.device)
            self.model.eval()

            # 使用 torch.compile 优化模型（类似 LiteRT 的编译优化）
            # 注意：某些操作可能不支持编译，会产生警告但不影响功能
            # 可以通过环境变量 ENABLE_TORCH_COMPILE=false 禁用编译
            enable_compile = os.getenv("ENABLE_TORCH_COMPILE", "true").lower() == "true"
            
            if enable_compile:
                logger.info("尝试编译模型以启用优化...")
                try:
                    # PyTorch 2.0+ 的 torch.compile 优化
                    if hasattr(torch, 'compile'):
                        # CPU 模式下使用 "default" 模式，避免 cudagraph 警告
                        # "reduce-overhead" 模式会尝试使用 cudagraph，在 CPU 上会产生警告
                        import warnings
                        with warnings.catch_warnings():
                            warnings.filterwarnings("ignore", category=UserWarning, message=".*cudagraph.*")
                            self._compiled_model = torch.compile(
                                self.model,
                                mode="default",  # CPU 上使用 default 模式，避免 cudagraph
                                fullgraph=False,
                                dynamic=False  # 禁用动态形状，减少警告
                            )
                        logger.info("✓ 模型 torch.compile 优化成功（CPU模式）")
                    else:
                        logger.warning("PyTorch 版本不支持 torch.compile，使用标准模式")
                        self._compiled_model = None
                except Exception as e:
                    logger.warning(f"编译优化警告: {e}，将使用标准 PyTorch 执行")
                    self._compiled_model = None
            else:
                logger.info("torch.compile 已禁用（通过环境变量），使用标准模式")
                self._compiled_model = None

            logger.info("✓ Gemma-3-270M 模型加载成功（PyTorch + 编译优化）")
            return True

        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}")
            raise

    def get_model(self):
        """获取模型（优先返回编译后的模型）"""
        if self.model is None:
            raise RuntimeError("模型未加载，请先调用load_model()")
        # 如果编译成功，返回编译后的模型，否则返回原始模型
        return self._compiled_model if self._compiled_model is not None else self.model

    def get_tokenizer(self):
        """获取tokenizer"""
        if self.tokenizer is None:
            raise RuntimeError("Tokenizer未加载，请先调用load_model()")
        return self.tokenizer

    def get_device(self):
        """获取设备"""
        return self.device
